#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <complex.h>
#include <string.h>
#include <unistd.h>

#include <sys/types.h> 
#include <sys/stat.h> 
#include <fcntl.h>

#include <gsl/gsl_complex.h>
#include <gsl/gsl_complex_math.h>
#include <gsl/gsl_vector_complex.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_fft_complex.h>
#include <gsl/gsl_matrix.h>

typedef struct
{
	int N;
	int M;
	int R;
	int delzeta;
	double Dtau;
	double taumin;
	int taumax;
	double a;
	double lambda;
	double zetamin;
	double zetamax;
	double zeta0;
	double Dzeta;
	double gama;
	double qmin;
	double Dq;
} info;

void imprimeReal(double* zeta, double* atual, int N, int m);
void imprimeTrans(double* atualK, info p, int m);
void initInfo(info* p);
void setInitialCondition(double* atual, double* Z, info p);
void nonLinear(double* nlinear, double* atual, double* zeta, info p);
void printPhaseReal(double* zeta, double* atual, int N, int m);
void printPhaseTrans(double* atualK, info p, int m);
void printFieldTau(gsl_matrix* fieldtau, info p);
void printfFieldOmega(double* fieldomega, info p, int i);

/*
	************************************************

	Start The physical and computational parameters of the system

	************************************************
*/

void initInfo(info* p)
{
	p->M = 16384; 			// 8192 values of tau per unit of tau
	p->N = 2048; 			// 1024 values of chi per unit of chi
	p->R = 32; 				// 16 values of q (p->N/64)
	p->delzeta = 16;
	p->zetamin = 0;
	p->zetamax = 2*p->N/p->delzeta;
	// Dzeta = (zetamax - zetamin)/N;
	p->Dzeta = 2.0/p->delzeta;
	p->taumin = 0;
	p->taumax = 128;
	p->Dtau = 1.0/p->M;
	p->a = 1.0;
	p->lambda = 1.0;
	p->gama = 10.0;
	p->zeta0 = p->N/(2.0*p->delzeta);
	// qmin = -1/2*Dzeta 
	p->qmin = -p->delzeta/4.0;
	p->Dq = 1.0*p->delzeta/(2.0*p->N);
}


/*
	************************************************

	The complex Electric Field must have the real and imaginary parts in alternating order, to use GSL Fourier Transform. It
	is called there compact notation.

	************************************************
*/

void setInitialCondition(double* atual, double* zeta, info p)
{
	int i;
	// double q = 0.2;
	
	for ( i = 0 ; i < p.N ; i++ )
	{
		zeta[i] = (p.zetamin + i*p.Dzeta);
		atual[2*i] = exp((-1.0*(zeta[i]-p.zeta0)*(zeta[i]-p.zeta0))/(2.0*p.gama*p.gama)); /* Parte Real */
		atual[2*i+1] = 0; /* Parte Imaginária */
	}
} 

/*
	************************************************
	
	The nonlinear term, using a saturable nonlinearity or a Kerr nonlinearity. It is the product of the
	refraction index(dependent on zeta and E) with the electric Field.
	
	n=n_{0}+n_{NL}\left|E\right|^{2}
	
	n_{0} = n_{1} or n_{2}, depending on \xi (famous zeta)
	
	nlinear=(n^{2}-1)E
	
	n_1 = 1
	n_2 = 1.25
	n_NL = 0;
	
	************************************************
*/

void nonLinear(double* nlinear, double* atual, double* zeta, info p)
{
	int i;
	double zetaBarrier = p.zetamax/2.0;
	int tmp;
	double n1, n2, nNL;
	double length = 20.0; // total length
	double a = 1.0; // Layer size
	
	n1 = 1.0;
	n2 = 1.25;
	nNL = 0.0;
	
	for ( i = 0 ; i < p.N ; i++ )
	{
		if ( zeta[i] < zetaBarrier )
		{
			nlinear[2*i] = 0.0; // Real Part equals zero
			nlinear[2*i+1] = 0.0; // Imag Part equals zero
			continue;
		}
// 		else
// 		{
// 			nlinear[2*i] = (n2*n2 - 1.0)*atual[2*i];
// 			nlinear[2*i+1] = (n2*n2 - 1.0)*atual[2*i+1];
// 			continue;
// 		}
		if ( fabs(zeta[i]-zetaBarrier) >= length )
		{
			nlinear[2*i] = 0.0; // Real Part equals zero
			nlinear[2*i+1] = 0.0; // Imag Part equals zero
			continue;
		}
		tmp = (zeta[i]-zetaBarrier)/a;
		if ( tmp%2 == 0 )
		{
			// Camada Par(even), linear
			nlinear[2*i] = 0.0;
			nlinear[2*i+1] = 0.0;
		}
		else
		{
			// Camada Impar(odd), nonlinear
			nlinear[2*i] = (n2*n2 - 1.0)*atual[2*i];
			nlinear[2*i+1] = (n2*n2 - 1.0)*atual[2*i+1];
		}
	}
}


int main(void)
{
//	double kb;
	int i, l, m, j, c1;
	double q;
	double *atual, *atualK, *proxK, *prox, *nlinearK, *nlinear, *zeta, *fieldomega;
	gsl_matrix* fieldtau;
	info p;
	double tau;
	// int fd;
	
	/* Define some initial conditions for the problem */
	initInfo(&p);
	// fd = open("./resultado/field.bin", O_WRONLY|O_CREAT, S_IWUSR|S_IRUSR);
	
	/* Alocate memory */
	atual = malloc(2*p.N*sizeof(double)); /* The electric field at time \tau */
	atualK = malloc(2*p.N*sizeof(double)); /* Fourier Transform of the Field */
	zeta = malloc(p.N*sizeof(double)); /* Spatial coordinate */
	prox = malloc(2*p.N*sizeof(double)); /* The electric Field at time \tau+1 */
	proxK = malloc(2*p.N*sizeof(double)); /* The Fourier Transform of the Electric Field at time \tau+1 */
	nlinear = malloc(2*p.N*sizeof(double)); /* nonlinear term */
	nlinearK = malloc(2*p.N*sizeof(double)); /* Fourier Transform of the nonlinear term */
	fieldtau = gsl_matrix_alloc(p.taumax*p.M, 2*(p.N/p.R)+1); /* The Field, as a function of time, for some q values */
	fieldomega = malloc(2*p.taumax*p.M*sizeof(double)); /* The Fourier Transform of the field, for fixed q. Gives us omega */ 
	
	/* Initial Value of the Electric Field and zeta vector */
	setInitialCondition(atual, zeta, p);  

	// write(fd, &p, sizeof(info));
	
	m = 0;

	/* loop start */
	for ( l = 0 ; l < p.taumax*p.M ; l++ )
	{
		tau = p.taumin + l*p.Dtau;
		if ( l == 0 )
		{
			/* Copy from atual to atualK, and obtain the Fourier Transform in atualK */
			memcpy(atualK, atual, 2*p.N*sizeof(double));
			gsl_fft_complex_radix2_forward(atualK, 1, p.N);
			
			memcpy(prox, atual, 2*p.N*sizeof(double));
			memcpy(proxK, atualK, 2*p.N*sizeof(double));
			
//			write(fd, prox, 2*p.N*sizeof(double));
			imprimeReal(zeta, prox, p.N, m);
			imprimeTrans(proxK, p, m);
			printPhaseReal(zeta, prox, p.N, m);
			printPhaseTrans(proxK, p, m);
			
			for (i = 0 ; i < p.N/2 ; i++ )
			{
				j = i + p.N/2;
				if ( j % p.R == 0 )
				{
					c1 = j/p.R;
					gsl_matrix_set(fieldtau, l, 0, tau);
					gsl_matrix_set(fieldtau, l, 2*c1+1, proxK[2*c1*p.R]);
					gsl_matrix_set(fieldtau, l, 2*c1+2, proxK[2*c1*p.R+1]);
				}
			}
			for ( i = p.N/2 ; i < p.N ; i++ )
			{
				j = i - p.N/2;
				if ( j % p.R == 0 )
				{
					c1 = j/p.R;
					gsl_matrix_set(fieldtau, l, 0, tau);
					gsl_matrix_set(fieldtau, l, 2*c1+1, proxK[2*c1*p.R]);
					gsl_matrix_set(fieldtau, l, 2*c1+2, proxK[2*c1*p.R+1]);
				}
			}
		}
		else
		{
			memcpy(atual, prox, 2*p.N*sizeof(double));
			memcpy(atualK, proxK, 2*p.N*sizeof(double));
			
			nonLinear(nlinear, atual, zeta, p);
			memcpy(nlinearK, nlinear, 2*p.N*sizeof(double));
			gsl_fft_complex_radix2_forward(nlinearK, 1, p.N);
			
			for ( i = 0 ; i < p.N/2 ; i++ )
			{
				j = i + p.N/2;
				q = (1.0*p.qmin + j*p.Dq)*2.0*M_PI; /* Coloca o fator 2\pi aqui para facilitar */
				proxK[2*i] = 	(((p.lambda*q*q*p.Dtau)/(4.0*M_PI*p.a)) + q*p.Dtau) * atualK[2*i+1]
								- M_PI*p.Dtau*nlinearK[2*i+1]
								+ atualK[2*i];
				proxK[2*i+1] = 	(-((p.lambda*q*q*p.Dtau)/(4.0*M_PI*p.a)) - q*p.Dtau) * atualK[2*i]
								+ M_PI*p.Dtau*nlinearK[2*i]
								+ atualK[2*i+1];
			}
			for ( i = p.N/2 ; i < p.N ; i++ )
			{
				j = i - p.N/2;
				q = (1.0*p.qmin + j*p.Dq)*2.0*M_PI;
				proxK[2*i] = 	(((p.lambda*q*q*p.Dtau)/(4.0*M_PI*p.a)) + q*p.Dtau) * atualK[2*i+1]
								- M_PI*p.Dtau*nlinearK[2*i+1]
								+ atualK[2*i];
				proxK[2*i+1] = 	(-((p.lambda*q*q*p.Dtau)/(4.0*M_PI*p.a)) - q*p.Dtau) * atualK[2*i]
								+ M_PI*p.Dtau*nlinearK[2*i]
								+ atualK[2*i+1];
				
			}
			
			memcpy(prox, proxK, 2*p.N*sizeof(double));
			gsl_fft_complex_radix2_inverse(prox, 1, p.N);

			// Save the complex values of the field, for all zeta and tau.
// 			write(fd, prox, 2*p.N*sizeof(double)); // Field is saved 
			
			// For some k values, save the field as a function of time. Later, we will calculate the Time-Fourier Transform
			for (i = 0 ; i < p.N/2 ; i++ )
			{
				j = i + p.N/2;
				if ( j % p.R == 0 )
				{
					c1 = j/p.R;
					gsl_matrix_set(fieldtau, l, 0, tau);
					gsl_matrix_set(fieldtau, l, 2*c1+1, proxK[2*c1*p.R]);
					gsl_matrix_set(fieldtau, l, 2*c1+2, proxK[2*c1*p.R+1]);
				}
			}
			for ( i = p.N/2 ; i < p.N ; i++ )
			{
				j = i - p.N/2;
				if ( j % p.R == 0 )
				{
					c1 = j/p.R;
					gsl_matrix_set(fieldtau, l, 0, tau);
					gsl_matrix_set(fieldtau, l, 2*c1+1, proxK[2*c1*p.R]);
					gsl_matrix_set(fieldtau, l, 2*c1+2, proxK[2*c1*p.R+1]);
				}
			}
			/* fieldtau has the right order, from lowest q to highest q */
			
			if ( l % p.M == 0 )
			{
				m++;
				imprimeReal(zeta, prox, p.N, m);
				printPhaseReal(zeta, prox, p.N, m);
				imprimeTrans(proxK, p, m);
				printPhaseTrans(proxK, p, m);
			}
		}
	}
	
	/* Now save fieldtau */
	printFieldTau(fieldtau, p);
	
	/* Now lets calculate the Fourier Transform, for some fixed q. It is the time-omega transform now. Calculate and save */
	for ( i = 0 ; i < p.N/p.R ; i++ )
	{
		/* Transform from gsl_matrix to compact notation, as required by gsl_fft */
		for ( l = 0 ; l < p.taumax*p.M ; l++ )
		{
			fieldomega[2*l] = gsl_matrix_get(fieldtau, l, 2*i+1);
			fieldomega[2*l+1] = gsl_matrix_get(fieldtau, l, 2*i+2);
		}
		gsl_fft_complex_radix2_forward(fieldomega, 1, p.taumax*p.M); /* Now fieldomega has the (time)Fourier Tranform, for fixed q */
		printfFieldOmega(fieldomega, p, i);
	}
	
//	close(fd);
	return 0;
}

void printfFieldOmega(double* fieldomega, info p, int i)
{
	int j, l;
	char str[100];
	FILE *fp;
	double w;
	
	if ( i < p.N/(2*p.R) )
	{
		j = i + p.N/(2*p.R);
	}
	else
	{
		j = i - p.N/(2*p.R);
	}
	
	sprintf(str, "resultado/arqW_Field_%d", j);
	fp = fopen(str, "w+");

	/* omega goes from -1/(2*deltatau) to 1/(2*deltatau) */
	/* deltatau = 1/M */
	
	for ( l = (p.taumax*p.M)/2 ; l < (p.taumax*p.M) ; l++ )
	{
		j = l - (p.taumax*p.M)/2;
		w = -p.M/2 + (1.0*j)/(1.0*p.taumax);
		fprintf (fp, "%14.6f\t%14.6f\n", w, (fieldomega[2*l]*fieldomega[2*l] + fieldomega[2*l+1]*fieldomega[2*l+1])/p.taumax*p.M);
	}
	for ( l = 0 ; l < (p.taumax*p.M)/2 ; l++ )
	{
		j = l + (p.taumax*p.M)/2;
		w = -p.M/2 + (1.0*j)/(1.0*p.taumax);
		fprintf (fp, "%14.6f\t%14.6f\n", w, (fieldomega[2*l]*fieldomega[2*l] + fieldomega[2*l+1]*fieldomega[2*l+1])/p.taumax*p.M);
	}

	fclose(fp);
}

void printFieldTau(gsl_matrix* fieldtau, info p)
{
	int i, l, j;
	char str[100];
	FILE *fp;
	
	
	for ( i = 0 ; i < p.N/p.R ; i++ )
	{
		if ( i < p.N/(2*p.R) )
		{
			j = i + p.N/(2*p.R);
		}
		else
		{
			j = i - p.N/(2*p.R);
		}
		sprintf(str, "resultado/arqT_Field_%d", j);
		fp = fopen(str, "w+");
		for ( l = 0 ; l < p.taumax*p.M ; l++ )
		{
			fprintf(fp, "%14.6f\t%14.6f\t%14.6f\n", gsl_matrix_get(fieldtau, l, 0), gsl_matrix_get(fieldtau, l, 2*i+1),
													gsl_matrix_get(fieldtau, l, 2*i+2));
		}
		fclose(fp);
	}
}

void imprimeReal(double* zeta, double* atual, int N, int m)
{
	int i;
	char str[100];
	FILE *fp;
	sprintf(str, "resultado/arqR_E2_%d", m);
	fp = fopen(str, "w+");
	
	for ( i = 0 ; i < N ; i++ )
	{
		fprintf(fp, "%14.6f\t%14.6f\n", zeta[i], atual[2*i]*atual[2*i]+atual[2*i+1]*atual[2*i+1]);
	}
	fclose(fp);
}

void printPhaseReal(double* zeta, double* atual, int N, int m)
{
	int i;
	char str[100];
	FILE *fp;
	double tmp;
	sprintf(str, "resultado/arqR_Phase_%d", m);
	fp = fopen(str, "w+");
	
	for ( i = 0 ; i < N ; i++ )
	{
		if ((atual[2*i]*atual[2*i]+atual[2*i+1]*atual[2*i+1]) < 1e-8)
			tmp = 0;
		else
			tmp = (atual[2*i+1]/atual[2*i]); 
		fprintf(fp, "%14.6f\t%14.6f\n", zeta[i], tmp);
	}
	fclose(fp);
}

void imprimeTrans(double* atualK, info p, int m)
{
	int i, j;
	char str[100];
	FILE *fp;
	double q;
	
	sprintf(str, "resultado/arqK_E2_%d", m);
	fp = fopen(str, "w+");
	
	/* Guardar os negativos primeiro, em ordem crescente */
	for ( i = p.N/2; i < p.N; i++)
	{
		j = i - p.N/2;
		q = (1.0*p.qmin + j*p.Dq);
		fprintf (fp, "%14.6f\t%14.6f\n", q, (atualK[2*i]*atualK[2*i] + atualK[2*i+1]*atualK[2*i+1])/p.N);
	}
	/* Agora vem os positivos, começando com o zero */
	for ( i = 0 ; i < p.N/2 ; i++ )
	{
		j = i + p.N/2;
		q = (1.0*p.qmin + j*p.Dq);
		fprintf (fp, "%14.6f\t%14.6f\n", q, (atualK[2*i]*atualK[2*i] + atualK[2*i+1]*atualK[2*i+1])/p.N);
	}
	fclose(fp);
}

void printPhaseTrans(double* atualK, info p, int m)
{
	int i, j;
	char str[100];
	FILE *fp;
	double q;
	double tmp;
	
	sprintf(str, "resultado/arqK_Phase_%d", m);
	fp = fopen(str, "w+");
	
	/* Guardar os negativos primeiro, em ordem crescente */
	for ( i = p.N/2; i < p.N; i++)
	{
		j = i - p.N/2;
		q = (1.0*p.qmin + j*p.Dq);
		if ((atualK[2*i]*atualK[2*i]+atualK[2*i+1]*atualK[2*i+1]) < 1)
			tmp = 0;
		else
			tmp = (atualK[2*i+1]/atualK[2*i]); 
		fprintf (fp, "%14.6f\t%14.6f\n", q, tmp);
	}
	/* Agora vem os positivos, começando com o zero */
	for ( i = 0 ; i < p.N/2 ; i++ )
	{
		j = i + p.N/2;
		q = (1.0*p.qmin + j*p.Dq);
		if ((atualK[2*i]*atualK[2*i]+atualK[2*i+1]*atualK[2*i+1]) < 1)
			tmp = 0;
		else
			tmp = (atualK[2*i+1]/atualK[2*i]); 
		fprintf (fp, "%14.6f\t%14.6f\n", q, tmp);
	}
	fclose(fp);
}

// gcc -Wall -lm -O3 -lgsl -lgslcblas main.c -o main.exe
